
import time
import threading
import subprocess
import math
from .gpu import gpus_snapshot, GPU

# %%
def compute_dicts_ema(new, cache, ema_coef):
  for idx, d in new.items():
    if not idx in cache:
      cache[idx] = d
      continue
    d_pre = cache.get(idx)
    for k, v in d.items():
      v_pre = d_pre.get(k)
      v_pre = v_pre if not v_pre is None else v
      if not v_pre is None and isinstance(v, float) and isinstance(v_pre, float) and not math.isnan(v_pre):
        d_pre[k] = ema_coef * v_pre + (1 - ema_coef) * v
      else:
        d_pre[k] = v
  return cache

# %%
class Monitor(threading.Thread):
  def __init__(self, interval=1, auto_start=True, ema_coef=0.6):
    super().__init__()
    self.interval = interval
    self._running = False
    self._running_lock = threading.Lock() 
    self._states_cache = {}
    self._states_lock = threading.Lock()
    self.ema_coef = ema_coef
    self._gpus = {}
    self.daemon = True
    self._refresh_gpu_status() # 构造时刷新一次
    if auto_start:
      self.start()
  
  def start(self):
    if self.is_alive():
      self.stop()
      self.join() 
    self._started._flag = False
    with self._running_lock:
      self._running = True
    super().start()
  
  def run(self):
    while self._running:
      self._refresh_gpu_status()
      time.sleep(self.interval)

  def stop(self):
    with self._running_lock:
      self._running = False

  def _refresh_gpu_status(self):
    try:
      ret, idxs = gpus_snapshot()
      with self._states_lock:
        newd = dict((idx, d) for idx,d in zip(idxs,ret))
        cache = self._states_cache.copy()
        self._states_cache = compute_dicts_ema(newd, cache, self.ema_coef)
      for idx, d in self._states_cache.items():
        if not idx in self._gpus:
          self._gpus[idx] = GPU(auto_update=False, **d)
        else:
          self._gpus[idx].update_from_dict(d)
    except Exception as e:
      print("Exception occurred:", e)
      
  @property
  def gpus(self):
    return list(self._gpus.values())
  
  @property
  def is_blank(self):
    if not self._states_cache:
      return True
    return False
      
  def print_gpus_status(self):
    # 打印表头
    print(f"{'ID':>3}  {'TEMP(C)':>4}  {'POWER(%)':>4}  {'LOAD(%)':>4}  {'MEM(%)':>4}")
    print("-" * 20)  # 打印分隔线
    for gpu in self.gpus:
      print(gpu.format)
      
  def filter(self, th_load=1.0, th_mem=0, th_pmem=1.0, th_pow=300, th_ppow=1.0, th_temp=80):
    if self.is_blank:
      self._refresh_gpu_status()
    ret = []
    self.gpus: list[GPU] # pyright: ignore
    for gpu in self.gpus:
      if not (gpu.load < th_load and gpu.power_percent < th_ppow and gpu.memory_percent < th_pmem):
        continue
      if not (gpu.temperature <= th_temp and gpu.power <= th_pow and gpu.memory_free >= th_mem):
        continue
      ret.append(gpu)
    return ret
  
  # 序列化时仅拷贝信息副本，不拷贝线程相关内容
  def __getstate__(self):
      state = {}
      state["_gpus"] = self.__dict__.get("_gpus", None)
      state["_states_cache"] = self.__dict__.get("_states_cache", None)
      return state